import os
from paddle.io import DataLoader
import paddle.vision.transforms as transforms

import load_data.mytransforms as mytransforms
from load_data.constant import tusimple_row_anchor, culane_row_anchor
from load_data.dataset import LaneClsDataset, LaneTestDataset

def get_train_loader(batch_size, data_root, griding_num, dataset, use_aux, num_lanes):
    target_transform = transforms.Compose([
        # mytransforms.FreeScaleMask((288, 800)),
        mytransforms.FreeScaleMask((560, 1000)),
        mytransforms.MaskToTensor(),
    ])
    segment_transform = transforms.Compose([
        mytransforms.FreeScaleMask((36, 100)),
        mytransforms.MaskToTensor(),
    ])
    img_transform = transforms.Compose([
        # transforms.Resize((288, 800)),
        transforms.Resize((560, 1000)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    simu_transform = mytransforms.Compose2([
        mytransforms.RandomRotate(6),
        mytransforms.RandomUDoffsetLABEL(100),
        mytransforms.RandomLROffsetLABEL(200)
    ])
    if dataset == 'CULane':
        train_dataset = LaneClsDataset(data_root,
                                           os.path.join(data_root, 'list/train_gt.txt'),
                                           img_transform=img_transform, target_transform=target_transform,
                                           simu_transform = simu_transform,
                                           segment_transform=segment_transform, 
                                           row_anchor = culane_row_anchor,
                                           griding_num=griding_num, use_aux=use_aux, num_lanes = num_lanes)
        cls_num_per_lane = 18

    elif dataset == 'TuSimple':
        train_dataset = LaneClsDataset(data_root,
                                           os.path.join(data_root, 'train_gt.txt'),
                                           img_transform=img_transform, 
                                           target_transform=target_transform,
                                           simu_transform = simu_transform,
                                           griding_num=griding_num, 
                                           row_anchor = tusimple_row_anchor,
                                           segment_transform=segment_transform,
                                           use_aux=use_aux, 
                                           num_lanes = num_lanes)
        cls_num_per_lane = 56
    else:
        raise NotImplementedError

    train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=0)

    return train_dataset, train_loader, cls_num_per_lane

def get_test_loader(batch_size, data_root,dataset):
    img_transforms = transforms.Compose([
        transforms.Resize((288, 800)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    if dataset == 'CULane':
        test_dataset = LaneTestDataset(data_root,os.path.join(data_root, 'list/test.txt'),img_transform = img_transforms)
        cls_num_per_lane = 18
    elif dataset == 'Tusimple':
        test_dataset = LaneTestDataset(data_root,os.path.join(data_root, 'test.txt'), img_transform = img_transforms)
        cls_num_per_lane = 56
    loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=4)
    return loader